import time

from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional

# 4. 定义请求体结构
from ai_quickstart.openai_translator.ai_translator.model import OpenAIModel
from ai_quickstart.openai_translator.ai_translator.translator import PDFTranslator
from ai_quickstart.openai_translator.ai_translator.utils import ConfigLoader, ArgumentParser


class APIService:
    def __init__(self):
        self.app = FastAPI()
        # 设置 CORS 中间件
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],  # 或者指定你允许的域名列表
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        args = ArgumentParser().parse_arguments()
        args.config = '../../config.yaml'
        # 读取配置文件
        config = ConfigLoader(args.config).load_config()

        # 从参数，或配置文件中获取执行所需的参数
        model_name = args.openai_model if args.openai_model else config['OpenAIModel']['model']
        api_key = args.openai_api_key if args.openai_api_key else config['OpenAIModel']['api_key']
        base_url = args.openai_base_url if args.openai_base_url else config['OpenAIModel']['base_url']
        # 封装OpenAI客户端
        model = OpenAIModel(model=model_name, api_key=api_key, base_url=base_url)
        self.translator = PDFTranslator(model)

    def setup_routes(self):
        @self.app.post("/upload/")
        def upload(file: UploadFile):
            # 文件对象属性，文件名，文件大小
            print(file.filename, file.size)
            return "OK"

        @self.app.post("/translate_pdf/")
        async def translate_pdf(file: UploadFile, file_format: str = Form("pdf"),
                                target_language: str = Form("中文"), pages: Optional[int] = Form(None)):
            try:
                print("Received request:",
                      {"file_format": file_format, "target_language": target_language, "pages": pages})  # 打印请求体内容，帮助调试

                str_time = time.strftime('%Y%m%d_%H%M%S', time.localtime(time.time()))
                temp_pdf_path = f"../../data/tmp/{str_time}_{file.filename}"
                output_file_file = f'{str_time}_' + file.filename.replace(".pdf", f"_translated.{file_format}")
                output_file_path = f"../../data/tmp/{output_file_file}"

                with open(temp_pdf_path, "wb") as temp_file:
                    temp_file.write(file.file.read())

                # 实例化 PDFTranslator 类，并调用 translate_pdf() 方法
                self.translator.translate_pdf(temp_pdf_path, file_format=file_format, output_file_path=output_file_path,
                                              target_language=target_language, pages=pages)

                return {"download_filename": output_file_file, "format": file_format}
                # return JSONResponse(content={"download_filename": output_file_file, "format": file_format})

            except Exception as e:
                print(f"Error: {str(e)}")
                return {"error": f"Internal server error: {str(e)}"}

        @self.app.get("/download/{filename}")
        async def download_file(filename: str):
            file_path = f"../../data/tmp/{filename}"
            return FileResponse(file_path)


if __name__ == '__main__':
    import uvicorn

    app_service = APIService()
    app_service.setup_routes()  # 注册路由

    uvicorn.run(app_service.app, host="127.0.0.1", port=8001)
